using System;
using System.IO;
using Org.BouncyCastle.Bcpg;
using Org.BouncyCastle.Bcpg.OpenPgp;
using Org.BouncyCastle.Security;
using Org.BouncyCastle.Utilities.IO;

namespace ICommonLibrary.Security {

    public class Pgp {

        /// <summary>
        /// Decrypt the passed in message stream
        /// </summary>
        /// <param name="input"></param>
        /// <param name="secretKeys"></param>
        /// <param name="passwd"></param>
        public static byte[] Decrypt(Stream input, PgpSecretKeyRingBundle secretKeys, char[] passwd) {

            input = PgpUtilities.GetDecoderStream(input);

            PgpObject obj = new PgpObjectFactory(input).NextPgpObject();
            
            // the first object might be a PGP marker packet.
            PgpEncryptedDataList objList;
            if (obj is PgpEncryptedDataList) {
                objList = (PgpEncryptedDataList)obj;
            }
            else {
                objList = (PgpEncryptedDataList)new PgpObjectFactory(input).NextPgpObject();
            }

            // find the secret key
            PgpPrivateKey secretKey = null;
            PgpPublicKeyEncryptedData encryptedData = null;

            foreach (PgpPublicKeyEncryptedData next in objList.GetEncryptedDataObjects()) {
                secretKey = FindSecretKey(secretKeys, next.KeyId, passwd);
                if (secretKey != null) {
                    encryptedData = next;
                    break;
                }
            }

            if (secretKey == null) {
                throw new ArgumentException("secret key for message not found.");
            }

            PgpObject message = new PgpObjectFactory(encryptedData.GetDataStream(secretKey)).NextPgpObject();

            if (message is PgpCompressedData) {
                var compressedMessage = (PgpCompressedData)message;
                message = new PgpObjectFactory(compressedMessage.GetDataStream()).NextPgpObject();
            }

            if (message is PgpLiteralData) {
                var pgpData = (PgpLiteralData) message;
                var output = new MemoryStream();
                Streams.PipeAll(pgpData.GetInputStream(), output);
                return output.ToArray();
            }
            
            if (message is PgpOnePassSignatureList) {
                throw new PgpException("encrypted message contains a signed message - not literal data.");
            }
            throw new PgpException("message is not a simple encrypted file - type unknown.");
        }

        public static byte[] Encrypt(Stream inputStream, PgpPublicKey publicKey, string outputFileName) {
            var buffer = new MemoryStream();

            Stream outputStream = buffer;
            const int BUFFER_SIZE = 1 << 16; // should always be power of 2
            bool armor = true;
            bool withIntegrityCheck = true;

            if (armor)
                outputStream = new ArmoredOutputStream(outputStream);

            // Init encrypted data generator
            var encryptedDataGenerator =
                new PgpEncryptedDataGenerator(SymmetricKeyAlgorithmTag.Cast5, withIntegrityCheck, new SecureRandom());
            encryptedDataGenerator.AddMethod(publicKey);
            Stream encryptedOut = encryptedDataGenerator.Open(outputStream, new byte[BUFFER_SIZE]);

            // Init compression
            var compressedDataGenerator = new PgpCompressedDataGenerator(CompressionAlgorithmTag.Zip);
            Stream compressedOut = compressedDataGenerator.Open(encryptedOut);

            // Create the Literal Data generator output stream
            var literalDataGenerator = new PgpLiteralDataGenerator();

            // TODO: Use lastwritetime from source file
            Stream literalOut = literalDataGenerator.Open(compressedOut, PgpLiteralData.Binary,
                outputFileName, DateTime.Now, new byte[BUFFER_SIZE]);

            // Open the input file
            var buf = new byte[BUFFER_SIZE];
            int len;
            while ((len = inputStream.Read(buf, 0, buf.Length)) > 0) {
                literalOut.Write(buf, 0, len);
            }

            literalOut.Close();
            literalDataGenerator.Close();

            compressedOut.Close();
            compressedDataGenerator.Close();
            encryptedOut.Close();
            encryptedDataGenerator.Close();
            inputStream.Close();

            if (armor)
                outputStream.Close();

            return buffer.ToArray();
        }


        /// <summary>
        /// A simple routine that opens a key ring file and loads the first available key suitable for encryption.
        /// </summary>
        /// <param name="keyStream"></param>
        /// <returns></returns>
        public static PgpPublicKey FindPublicKey(Stream keyStream) {
            var keyRingBundle = new PgpPublicKeyRingBundle(PgpUtilities.GetDecoderStream(keyStream));

            foreach (PgpPublicKeyRing keyRing in keyRingBundle.GetKeyRings()) {
                foreach (PgpPublicKey key in keyRing.GetPublicKeys()) {
                    if (key.IsEncryptionKey) {
                        return key;
                    }
                }
            }
            throw new ArgumentException("Can't find encryption key in key ring.");
        }

        /// <summary>
        /// Search a secret key ring collection for a secret key corresponding to keyId if it exists.
        /// </summary>
        /// <param name="keyRingBundle"></param>
        /// <param name="keyId"></param>
        /// <param name="pass"></param>
        /// <returns></returns>
        public static PgpPrivateKey FindSecretKey(PgpSecretKeyRingBundle keyRingBundle, long keyId, char[] pass) {
            PgpSecretKey key = keyRingBundle.GetSecretKey(keyId);
            if (key == null) {
                return null;
            }
            return key.ExtractPrivateKey(pass);
        }

    }
}